module Zeno.HaskellParser 
(
  parseHaskellFile
)
where

import Prelude ()
import StdImports
import qualified Data.Map as Map
import System.Directory

import Zeno.Core
import Zeno.Term
import Zeno.Type
import Zeno.Function

import qualified GHC.Paths as Paths
import qualified CoreSyn as Hs
import qualified Var as Hs
import qualified Type as Hs
import qualified HscTypes as Hs
import qualified DataCon as Hs
import qualified TyCon as Hs
import qualified GHC as Hs
import qualified DynFlags as Hs
import qualified Module as Hs

import Outputable ( Outputable, ppr, defaultUserStyle )


data HsEnv = HsEnv  { envTerms :: Map (Either String Hs.Var) ZTerm,
                      envTypes :: Map (Either String Hs.Type) ZType }

emptyEnv = HsEnv    { envTerms = mempty,
                      envTypes = mempty }
                      
type HsZeno = StateT HsEnv ZenoState
type HsExpr = Hs.Expr Hs.Var
type HsBind = Hs.Bind Hs.Var
type HsBinding = (Hs.Var, HsExpr)

instance Eq Hs.Type where
  (==) = Hs.tcEqType
                     
instance Ord Hs.Type where
  compare = Hs.tcCmpType
  
instance WithinTraversable (Hs.Expr a) (Hs.Expr a) where
  mapWithinM f (Hs.App x y) = 
    f =<< Hs.App `liftM` mapWithinM f x `ap` mapWithinM f y 
  mapWithinM f (Hs.Let bind expr) = 
    f =<< Hs.Let `liftM` mapWithinM f bind `ap` mapWithinM f expr
  mapWithinM f (Hs.Lam x expr) = 
    f =<< Hs.Lam x `liftM` mapWithinM f expr
  mapWithinM f (Hs.Cast expr coer) = 
    f =<< (flip Hs.Cast coer) `liftM` mapWithinM f expr
  mapWithinM f (Hs.Note note expr) = 
    f =<< Hs.Note note `liftM` mapWithinM f expr
  mapWithinM f (Hs.Case expr b t alts) =
    f =<< Hs.Case `liftM` mapWithinM f expr `ap` 
      return b `ap` return t `ap` mapM traverseAlt alts
    where
      traverseAlt (con, b, expr) = do
        expr' <- mapWithinM f expr
        return (con, b, expr')      
  mapWithinM f x = f x
  
instance WithinTraversable (Hs.Expr a) (Hs.Bind a) where
  mapWithinM f (Hs.NonRec b ex) = Hs.NonRec b `liftM` mapWithinM f ex
  mapWithinM f (Hs.Rec binds) = Hs.Rec `liftM` mapM mapBinding binds
    where 
      mapBinding (b, ex) = do
        ex' <- mapWithinM f ex
        return (b, ex')
        
parseHaskellFile :: FilePath -> IO (ZenoState ())
parseHaskellFile filepath = do
  core_mod <- Hs.defaultErrorHandler Hs.defaultDynFlags $ 
    Hs.runGhc (Just Paths.libdir) $ do
      dflags <- Hs.getSessionDynFlags
      let dflags' = dflags {-
            { Hs.optLevel = 0,
              Hs.simplPhases = 2,
              Hs.maxSimplIterations = 6 } -}
      Hs.setSessionDynFlags $ dflags'
      Hs.compileToCoreSimplified filepath 
      
  when debug $ putStrLn (output core_mod)

  let types = Hs.cm_types core_mod
      binds = Hs.flattenBinds (Hs.cm_binds core_mod)

  return $ do
    flip evalStateT emptyEnv $ do
      --createBoolType
      createListType
      forM_ [2..5] createTupleType
      addEnvTerm (Left "Control.Exception.Base.patError") Err
      convertDataCons (Hs.typeEnvDataCons types)
      convertGlobalBindings binds
    
createBoolType :: HsZeno ()
createBoolType = do
  bool_id <- lift newIdS
  true_id <- lift newIdS
  false_id <- lift newIdS
  let true_var = ZVar true_id bool_type ZConstructor
      false_var = ZVar false_id bool_type ZConstructor
      bool_dtype = ZDataType bool_id [] [true_var, false_var]
      bool_type = TypeVar bool_dtype 
  addEnvVar  (Left "GHC.Bool.True")  true_var
  addEnvVar  (Left "GHC.Bool.False") false_var
  addEnvType (Left "GHC.Bool.Bool")  bool_dtype
  lift $ do christen bool_id "Bool" 
            christen true_id "True" 
            christen false_id "False" 
            addDataType bool_dtype
      
createListType :: HsZeno ()
createListType = do
  list_id <- lift newIdS
  empty_id <- lift newIdS 
  cons_id <- lift newIdS
  poly_id <- lift newIdS
  let poly_dtype = ZPolyType poly_id
      poly_type = TypeVar poly_dtype
      empty_type = AppType (TypeVar list_dtype) poly_type
      cons_type = FunType poly_type (FunType empty_type empty_type)
      list_dtype = ZDataType list_id [poly_dtype] [empty_var, cons_var]
      empty_var = ZVar empty_id empty_type ZConstructor
      cons_var = ZVar cons_id cons_type ZConstructor
  addEnvVar  (Left "GHC.Types.[]") empty_var
  addEnvVar  (Left "GHC.Types.:")  cons_var
  addEnvType (Left "[]") list_dtype
  lift $ do christen list_id "[]"
            christen empty_id "[]"
            christen cons_id ":"
            addDataType list_dtype
            
createTupleType :: Int -> HsZeno ()
createTupleType n = do
  poly_ids <- replicateM n (lift newIdS)
  con_id <- lift newIdS
  dtype_id <- lift newIdS
  let poly_dtypes = map ZPolyType poly_ids
      poly_types = map TypeVar poly_dtypes
      dtype = ZDataType dtype_id poly_dtypes [con_var]
      ret_type = unflattenAppType ((TypeVar dtype) : poly_types)
      con_type = unflattenFunType (poly_types ++ [ret_type])
      con_var = ZVar con_id con_type ZConstructor
      symbol = "(" ++ replicate (n - 1) ',' ++ ")"
  addEnvVar (Left $ "GHC.Tuple." ++ symbol) con_var
  addEnvType (Left symbol) dtype
  lift $ do christen con_id symbol
            christen dtype_id symbol
            addDataType dtype
            
addEnvTerm :: Either String Hs.Var -> ZTerm -> HsZeno ()
addEnvTerm var zterm = modify $ \env -> 
  env { envTerms = Map.insert var zterm (envTerms env) }
  
addEnvVar var = addEnvTerm var . Var
  
resultHsType :: Hs.Type -> Hs.Type
resultHsType t = 
  case Hs.splitFunTy_maybe t of
    Nothing -> t
    Just (_, t') -> resultHsType t'

createHsEnvVar :: Hs.Var -> ZVarClass -> HsZeno ZVar
createHsEnvVar hs_var var_cls = do
  new_id <- lift newIdS
  var_type <- convertCoreType (Hs.varType hs_var)
  let new_var = ZVar new_id var_type var_cls
  addEnvVar (Right hs_var) new_var
  lift $ christen new_id (outputName hs_var) 
  return new_var
  
addEnvType :: Either String Hs.Type -> ZTypeVar -> HsZeno ()
addEnvType type_name dtype = modify $ \env -> 
  env { envTypes = Map.insert type_name (TypeVar dtype) (envTypes env) }
  
output :: Outputable a => a -> String
output = show . flip ppr defaultUserStyle 

outputName :: Outputable a => a -> String
outputName = stripEndNumber . stripModuleName . output
  where 
    stripModuleName name = 
      if '.' `elem` name 
        then tail $ dropWhile (/= '.') name
        else name
    
    stripEndNumber name = 
      if isAlpha $ head name 
        then name
        else takeWhile (not . isDigit) name

lookupTypeId :: Hs.Type -> HsZeno ZType
lookupTypeId hs_type = do
  types <- envTypes <$> get
  case Map.lookup (Right hs_type) types of
    Just ztype -> return ztype
    Nothing -> 
      case Map.lookup (Left $ output hs_type) types of
        Just ztype -> return ztype
        Nothing -> traceZ ("Cannot find type " ++ output hs_type) $
          return EmptyType
    
lookupVar :: Hs.Var -> HsZeno ZTerm 
lookupVar hs_var = do
  terms <- envTerms <$> get
  case Map.lookup (Right hs_var) terms of
    Just zterm -> return zterm
    Nothing -> 
      case Map.lookup (Left $ output hs_var) terms of
        Just zterm -> return zterm
        Nothing -> traceZ ("Cannot find variable " ++ output hs_var) $
          return Err

convertCoreType :: Hs.Type -> HsZeno ZType
convertCoreType hs_type = do      
  let app_type = case Hs.splitAppTy_maybe hs_type of
        Just (fun, arg) -> do
          fun' <- convertCoreType fun
          arg' <- convertCoreType arg
          return (AppType fun' arg')
        Nothing -> 
          lookupTypeId hs_type
      
      fun_type = case Hs.splitFunTy_maybe hs_type of
        Just (arg, res) -> do
          arg' <- convertCoreType arg
          res' <- convertCoreType res
          return (FunType arg' res')
        Nothing -> 
          app_type
      
  case Hs.splitForAllTy_maybe hs_type of
    Just (tyvar, rest) -> do
      new_id <- lift newIdS
      addEnvType (Right $ Hs.mkTyVarTy tyvar) (ZPolyType new_id)
      convertCoreType rest
    Nothing ->
      fun_type  
      
convertGlobalBindings :: [HsBinding] -> HsZeno ()
convertGlobalBindings binds = do
  rec mapM_ (uncurry addEnvVar) binds'
      binds' <- mapM convertBinding binds
  return ()
  where
    convertBinding :: HsBinding -> HsZeno (Either String Hs.Var, ZVar)
    convertBinding (hs_var, expr) = do
      (args, stripped_expr) <- stripArgs expr
      fun_expr <- convertExpr (mapWithin transformExpr stripped_expr)
      fun_id <- lift newIdS
      fun_type <- convertCoreType (Hs.varType hs_var)
      let fun_var = ZVar fun_id fun_type (ZFunction fun)
          fun = Function fun_var args fun_expr
      lift $ do christen fun_id (outputName hs_var)
                addFunction fun
      return (Right hs_var, fun_var)

    stripArgs :: HsExpr -> HsZeno ([ZVar], HsExpr)
    stripArgs (Hs.Lam hs_var expr) 
      | Hs.isTyVar hs_var = do
          new_id <- lift newIdS
          addEnvType (Right $ Hs.mkTyVarTy hs_var) (ZPolyType new_id)
          stripArgs expr
      | otherwise = do
          z_var <- createHsEnvVar hs_var ZVariable
          addEnvVar (Right hs_var) z_var
          (vars, expr') <- stripArgs expr
          return (z_var : vars, expr')
    stripArgs expr = return ([], expr)
        
    transformExpr :: HsExpr -> HsExpr
    transformExpr (Hs.Cast expr _) = expr
    transformExpr (Hs.Note _ expr) = expr
    transformExpr (Hs.App (Hs.Lam b expr) expr_val) = 
      let replace (Hs.Var var) 
            | b == var = expr_val
          replace ex = ex
      in mapWithin transformExpr $ mapWithin replace expr
    transformExpr (Hs.Let (Hs.NonRec b expr1) expr2) = 
      let replace (Hs.Var var) 
            | b == var = expr1
          replace ex = ex
      in mapWithin transformExpr $ mapWithin replace expr2
    transformExpr (Hs.App expr (Hs.Type _)) = 
      transformExpr expr
    transformExpr (Hs.App x (Hs.Case of_expr hs_var ty alts)) = 
      Hs.Case of_expr hs_var ty $
        map (\(con, bs, expr) -> (con, bs, transformExpr (Hs.App x expr))) alts
    transformExpr (Hs.App (Hs.Case of_expr hs_var ty alts) y) = 
      Hs.Case of_expr hs_var ty $
        map (\(con, bs, expr) -> (con, bs, transformExpr (Hs.App expr y))) alts
      {-
    transformExpr expr@(Hs.Let (Hs.Rec _) _) = 
      error $ "Recursive let definition inside expression " ++ output expr 
    transformExpr expr@(Hs.Lam _ _) =
      error $ "Lambda abstraction inside expression " ++ output expr -}
    transformExpr expr = expr
        
    
convertTermExpr :: HsExpr -> HsZeno ZTerm
convertTermExpr (Hs.Var hs_var) = lookupVar hs_var
convertTermExpr (Hs.Lit _) = return Err
convertTermExpr (Hs.App fun arg) = do
  zfun <- convertTermExpr fun
  zarg <- convertTermExpr arg
  return $ App zfun zarg
convertTermExpr expr = traceZ ("Cannot convert expression " ++ output expr) $ 
  return Err

convertExpr :: HsExpr -> HsZeno ZExpression
convertExpr (Hs.Case of_expr hsvar _ alts) = do
  z_var <- createHsEnvVar hsvar ZVariable
  whens' <- mapM convertAlts alts
  of_term <- convertTermExpr of_expr
  let whens = map (replaceWithin (Var z_var) of_term) whens' 
  return $ CaseOf of_term whens
  where
    convertAlts :: Hs.Alt Hs.Var -> HsZeno ZWhen 
    convertAlts (Hs.DataAlt con, binds, expr) = do
      con_var <- lookupVar (Hs.dataConWorkId con)
      bind_vars <- mapM (flip createHsEnvVar ZVariable) binds
      z_expr <- convertExpr expr
      let term = unflattenTerm $ (con_var : (map Var bind_vars))
      return $ When term z_expr
convertExpr expr =
  Value <$> collapseErrors <$> convertTermExpr expr

  
convertDataCons :: [Hs.DataCon] -> HsZeno ()
convertDataCons cons = do
  let con_types = nub $ map dataConReturnType cons
      typed_cons = map (\t -> (t, consOfType t)) con_types
  
  rec mapM_ (uncurry addEnvType) types
      types <- mapM createDataType typed_cons
  return ()
  where
    createDataType :: (Hs.Type, [Hs.DataCon]) -> 
        HsZeno (Either String Hs.Type, ZTypeVar)
    createDataType (hs_type, cons) = do
      (hs_type', arg_types) <- stripTypeArgs hs_type
      type_id <- lift newIdS
      cons' <- mapM createDataConVar cons
      let datatype = ZDataType type_id arg_types cons'
      lift $ do addDataType datatype
                christen type_id (outputName hs_type')
      return (Right hs_type', datatype)
      where
        stripTypeArgs :: Hs.Type -> HsZeno (Hs.Type, [ZTypeVar])
        stripTypeArgs hs_type = do
          let (fun_type, arg_types) = Hs.splitAppTys hs_type
          ptypes <- mapM createPolyType arg_types
          return (fun_type, ptypes)
          where
            createPolyType :: Hs.Type -> HsZeno ZTypeVar
            createPolyType hs_type = do
              new_id <- lift newIdS
              lift $ christen new_id (outputName hs_type)
              let new_type = ZPolyType new_id
              addEnvType (Right hs_type) new_type
              return new_type
      
        -- No function types allowed yet
        createDataConVar :: Hs.DataCon -> HsZeno ZVar
        createDataConVar con = do
          mapM_ addTyVar (Hs.dataConAllTyVars con)
          var_id <- lift newIdS
          arg_types <- mapM convertCoreType (Hs.dataConOrigArgTys con)
          res_type <- convertCoreType (Hs.dataConOrigResTy con)
          let var_type = unflattenFunType (arg_types ++ [res_type])  
              new_var = ZVar var_id var_type ZConstructor
          lift $ christen var_id (outputName $ Hs.dataConName con)
          addEnvVar (Right $ Hs.dataConWorkId con) new_var
          return new_var
          where
            addTyVar :: Hs.TyVar -> HsZeno ()
            addTyVar tyvar = do
              type_id <- lift newIdS
              addEnvType (Right $ Hs.mkTyVarTy tyvar) (ZPolyType type_id)
                
    consOfType :: Hs.Type -> [Hs.DataCon]
    consOfType t = filter ((== t) . dataConReturnType) cons
  
    dataConReturnType :: Hs.DataCon -> Hs.Type
    dataConReturnType con = 
      let (_, _, _, t) = Hs.dataConSig con in t
      
      
      
      
